import subprocess, os
import time
import getpass
import yaml
import sys

utils_path = os.path.join('..', 'lib', 'utilities')
sys.path.append(utils_path)

import load_configurations
config_user_dic = load_configurations.open_yaml("../config_user.yaml")

# Specifying paths to shared directories for pipeline.
dropbox = str(config_user_dic["external"]["dropbox"])
shared_simulations = str(config_user_dic["external"]["shared_simulations"])
shared_estimations = str(config_user_dic["external"]["shared_estimations"])

full_run = str(config_user_dic["local"]["full_run"])
light     = config_user_dic["local"]["light"]
simulation     = config_user_dic["local"]["simulation"]

def Main():

    indir = 'code/'
    
    Estimate_jid = SubmitJob(indir, 'Estimate.sh', alt_results = full_run)

    if simulation:
        Rep_jid = SubmitJob(indir, 'Replicate.sh', alt_results = full_run)
        Ree_jid = SubmitJob(indir, 'Reestimate.sh', alt_results = full_run, dependency=Rep_jid)
        Sim_jid = SubmitJob(indir, 'Sim.sh', alt_results = full_run, dependency=Ree_jid)
        print('Jobs for replication, reestimation and simulation submitted!')
        

    if full_run:
        market_list = [10000]

    for market in market_list:
    
        custom_version_price_model_list = [
            ("", "mw", "L"),
            ("", "price", "L"),
            ("", "sg", "L"),
            ("", "coarse", "L"),
            ("", "prod", "L"),
            ("", "mw", "L_nm"),
            ("", "sg", "L_nm"),
            ("", "coarse", "L_nm"),
            ("", "mw", "RCL"),
            ("", "sg", "RCL"),
            ("_sh", "mw", "L"),
            ("_sh", "sg", "L"),
            ("_sh", "prod", "L")
        ]

        Est_jid_list = []
        Est_jid = cal_distance(indir, market, dependency = Sim_jid)
        Est_jid_list.append(Est_jid)   
                            
        for version, price_iv, model in custom_version_price_model_list:
            par_rest_list = get_lists(model, price_iv, market)
            for par_rest in par_rest_list:
                    Est_jid = SubmitEstimateJob(
                        indir, full_run,
                        market, price_iv, model, 
                        par_rest,
                        light,
                        dependency = Sim_jid, 
                        version = version
                    )

            Est_jid_list.append(Est_jid)

    dependency_list = ":".join(Est_jid_list)
    print('Jobs for estimation submitted!')

    Out_jid = SubmitJob(indir, 'OutputResults.sh', full_run, dependency = dependency_list)
    CSV_jid = SubmitJob(indir, 'ExportCSV.sh', full_run, dependency = Out_jid)

    print('Jobs for exporting results submitted!')

def SubmitJob(indir, job, alt_results, dependency = '',  markets = [10000]):

    dropbox_data_path = dropbox
    sim_data_path = shared_simulations
    job_file = indir + job
    
    if job == "Replicate.sh":

        f = open(job_file, "w")
        f.write("#!/bin/bash\n")
        f.write("#SBATCH -p gentzkow,hns,normal\n")
        f.write("#SBATCH --job-name=rep.job\n")
        f.write("#SBATCH --output=rep.out\n")
        f.write("#SBATCH --error=rep.err\n")
        f.write("#SBATCH --nodes=1\n")
        f.write("#SBATCH --ntasks=1\n")
        f.write("#SBATCH --cpus-per-task=1\n")
        f.write("#SBATCH --mem=8g\n")
        f.write("#SBATCH --time=05:00:00\n")
        f.write("#SBATCH --qos=normal\n")
        f.write("ml load matlab/R2022b\n")
        f.write("cd code\n")
        f.write(f"matlab -nodisplay -nosplash -nodesktop -r \"Replicate('{dropbox_data_path}', '{sim_data_path}', '{sim_data_path}'); exit;\" \n")
        f.write("cd -\n")
        f.close()

        if dependency:
            cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
        
        else:
            cmd = "sbatch --parsable %s" % job_file
    
        print("Submitting job for %s" % job)

        status, jid = subprocess.getstatusoutput(cmd)

        if (status == 0):
            print("Job id for %s is %s" % (job, jid))
        else:
            print("Error submitting job for %s" % job)
            print(subprocess.getstatusoutput(cmd))
    
    elif job == "Estimate.sh":
        
        f = open(job_file, "w")
        f.write("#!/bin/bash\n")
        f.write("#SBATCH -p gentzkow,hns,normal\n")
        f.write("#SBATCH --job-name=Estimate.job\n")
        f.write("#SBATCH --output=Estimate.out\n")
        f.write("#SBATCH --error=Estimate.err\n")
        f.write("#SBATCH --nodes=1\n")
        f.write("#SBATCH --ntasks=1\n")
        f.write("#SBATCH --cpus-per-task=1\n")
        f.write("#SBATCH --mem=8g\n")
        f.write("#SBATCH --time=05:00:00\n")
        f.write("#SBATCH --qos=normal\n")
        f.write("ml load matlab/R2022b\n")
        f.write("cd code\n")
        f.write(f"matlab -nodisplay -nosplash -nodesktop -r \"Estimate('{dropbox_data_path}', '{sim_data_path}', 'mw'); exit;\" \n")
        f.write(f"matlab -nodisplay -nosplash -nodesktop -r \"Estimate('{dropbox_data_path}', '{sim_data_path}', 'sg'); exit;\" \n")
        f.write("cd -\n")
        f.close()

        if dependency:
            cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
        
        else:
            cmd = "sbatch --parsable %s" % job_file
    
        print("Submitting job for %s" % job)

        status, jid = subprocess.getstatusoutput(cmd)

        if (status == 0):
            print("Job id for %s is %s" % (job, jid))
        else:
            print("Error submitting job for %s" % job)
            print(subprocess.getstatusoutput(cmd))

    elif job == "Reestimate.sh":
        
        f = open(job_file, "w")
        f.write("#!/bin/bash\n")
        f.write("#SBATCH -p gentzkow,hns,normal\n")
        f.write("#SBATCH --job-name=res.job\n")
        f.write("#SBATCH --output=%x_%a.out\n")
        f.write("#SBATCH --error=%x_%a.err\n")
        f.write("#SBATCH --nodes=1\n")
        f.write("#SBATCH --ntasks=1\n")
        f.write("#SBATCH --cpus-per-task=1\n")
        f.write("#SBATCH --mem=8g\n")
        f.write("#SBATCH --array=1-7\n")
        f.write("#SBATCH --time=04:00:00\n")
        f.write("#SBATCH --qos=normal\n")
        f.write("ml load matlab/R2022b\n")
        f.write("cd code\n")
        f.write(f"matlab -nodisplay -nosplash -nodesktop -r \"Reestimate('{dropbox_data_path}', '{sim_data_path}', '{sim_data_path}', $SLURM_ARRAY_TASK_ID, ''); exit;\" \n")
        f.write(f"matlab -nodisplay -nosplash -nodesktop -r \"Reestimate('{dropbox_data_path}', '{sim_data_path}', '{sim_data_path}', $SLURM_ARRAY_TASK_ID, '_sh'); exit;\" \n")
        f.write("cd -\n")
        f.close()

        if dependency:
            cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
        
        else:
            cmd = "sbatch --parsable %s" % job_file
    
        print("Submitting job for %s" % job)

        status, jid = subprocess.getstatusoutput(cmd)

        if (status == 0):
            print("Job id for %s is %s" % (job, jid))
        else:
            print("Error submitting job for %s" % job)
            print(subprocess.getstatusoutput(cmd))


    elif job == "Sim.sh":

        Sim_jid_list = []
        for market in markets:
            if light == False:
                arrayrange = range(1, 1101)
            if light == True:
                arrayrange = [i for start in range(1, 1002, 100) for i in range(start, start + 5)]
            for array in arrayrange:
                if array % 100 in (1, 2, 3):
                    jobtime = "80:00:00"
                    memory  = "20g"
                else:
                    jobtime = "24:00:00"
                    memory  = "8g"
                f = open(job_file, "w")
                f.write("#!/bin/bash\n")
                f.write("#SBATCH -p gentzkow,hns\n")
                f.write("#SBATCH --job-name=prod_%s.job\n" % (market))
                f.write("#SBATCH --output=prod_%s.out\n" % (market))
                f.write("#SBATCH --error=prod_%s.err\n" % (market))
                f.write("#SBATCH --nodes=1\n")
                f.write("#SBATCH --ntasks=1\n")
                f.write("#SBATCH --mem=%s\n" % (memory))
                f.write("#SBATCH --array=%s\n" % (array))
                f.write("#SBATCH --time=%s\n" % (jobtime))
                f.write("#SBATCH --qos=normal\n\n")
                f.write("ml load matlab/R2022b\n")
                f.write("cd code/Sim\n")
                f.write("START=$((1 + 1 * ($SLURM_ARRAY_TASK_ID - 1)))\n")
                f.write("STOP=$((0 + 1 * $SLURM_ARRAY_TASK_ID))\n")
                f.write("matlab -nodisplay -nosplash -nodesktop \\\n")
                f.write("""    -r "for s=$START:$STOP; Sim(s, '%s', '%s', '%s', %s, ''); end; \\\n""" % (dropbox_data_path, sim_data_path, sim_data_path, market))
                f.write("""        for s=$START:$STOP; Sim(s, '%s', '%s', '%s', %s, '_sh'); end; exit;"\n""" % (dropbox_data_path, sim_data_path, sim_data_path, market))
                f.write("cd -\n")

                f.close()
        
                if dependency:
                    cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
                else:
                    cmd = "sbatch --parsable %s" % job_file
                print("Submitting job for %s" % job_file)

                status, jid = subprocess.getstatusoutput(cmd)

                if (status == 0):
                    print("Job id for %s is %s" % (job, jid))
                else:
                    print("Error submitting job for %s" % job)
                    print(subprocess.getstatusoutput(cmd))
                
                Sim_jid_list.append(jid)

        jid = ":".join(Sim_jid_list)

    else:
            if dependency:
                cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
            else:
                cmd = "sbatch --parsable %s" % job_file
            print("Submitting job for %s" % job_file)

            status, jid = subprocess.getstatusoutput(cmd)

            if (status == 0):
                print("Job id for %s is %s" % (job, jid))
            else:
                print("Error submitting job for %s" % job)
                print(subprocess.getstatusoutput(cmd))

    return(jid)

def SubmitEstimateJob(indir, full_run, market, price_iv, model, par_rest, light,  dependency = '', version = ''):
    
    data_path = shared_simulations
    out_path  = shared_estimations
    user = getpass.getuser()

    jobname = str(market) + '_markets_' + price_iv[0] + price_iv[len(price_iv)-1] + "_" + model + "".join(str(p) for p in par_rest)

    if light == False:
        nodes, ntasks, nmem, time, array, arrayjob = task_allocation_full(model, par_rest, market)      
    if light == True:
        nodes, ntasks, nmem, time, array, arrayjob = task_allocation_light(model, market)

    job_file = os.path.join(indir, "%s.sh" % jobname)

    f = open(job_file, "w")
    f.write("#!/bin/bash\n")
    f.write("#SBATCH -p gentzkow,hns\n")
    f.write("#SBATCH --job-name=%s.job\n" % jobname)
    f.write("#SBATCH --output=%x_%a.out\n")
    f.write("#SBATCH --error=%x_%a.err\n")
    f.write("#SBATCH --nodes=%s\n" % nodes)
    f.write("#SBATCH --ntasks=%s\n" % ntasks)
    f.write("#SBATCH --mem=%sg\n" % nmem)
    f.write("#SBATCH --array=%s\n" % array)
    f.write("#SBATCH --time=%s\n" % time)
    f.write("#SBATCH --qos=normal\n")
    f.write("ml purge \n")
    f.write("ml system git \n")
    f.write("ml system git-lfs \n")
    f.write("ml curl/7.81.0 \n")
    f.write("ml system rclone \n")
    f.write("ml load matlab/R2022b \n")
    f.write("source /home/users/%s/miniconda3/bin/activate blp-instruments \n" %user)
    f.write("cd code/EstimateUnified \n")

    f.write('echo "Array Task ID: $SLURM_ARRAY_TASK_ID"\n')
    f.write(f"START=$((1 + {arrayjob} * ($SLURM_ARRAY_TASK_ID - 1)))\n")
    f.write(f"STOP=$((0 + {arrayjob} * $SLURM_ARRAY_TASK_ID))\n")
    f.write("matlab -nodisplay -nosplash -nodesktop \\\n")
    f.write("""    -r "for s=$START:$STOP; EstimateUnified('%s', '%s', %s, s, '%s', '%s', %s, '%s'); end; exit;"\n""" % (price_iv, model, par_rest, data_path, out_path, market, version))
    f.write("cd -")
    f.close()

    if dependency:
        cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
    else:
        cmd = "sbatch --parsable %s" % job_file

    print("Submitting job for %s" % jobname)
    status, jid = subprocess.getstatusoutput(cmd)

    if (status == 0):
        print("Job id for %s is %s" % (jobname, jid))
    else:
        print("Error submitting job for %s %s." % (jobname, status))
        print(subprocess.getstatusoutput(cmd))   

    return(jid)


def cal_distance(indir, market, dependency = ''):
    
    data_path = shared_simulations
    out_path  = shared_simulations
    user = getpass.getuser()

    jobname = 'Distance_' + str(market) + '_markets'

    nodes, ntasks, nmem, time, array, arrayjob = task_allocation_distance(market)

    job_file = os.path.join(indir, "%s.sh" % jobname)

    f = open(job_file, "w")
    f.write("#!/bin/bash\n")
    f.write("#SBATCH -p gentzkow,hns\n")
    f.write("#SBATCH --job-name=%s.job\n" % jobname)
    f.write("#SBATCH --output=%x_%a.out\n")
    f.write("#SBATCH --error=%x_%a.err\n")
    f.write("#SBATCH --nodes=%s\n" % nodes)
    f.write("#SBATCH --ntasks=%s\n" % ntasks)
    f.write("#SBATCH --mem=%sg\n" % nmem)
    f.write("#SBATCH --array=%s\n" % array)
    f.write("#SBATCH --time=%s\n" % time)
    f.write("#SBATCH --qos=normal\n")
    f.write("ml purge \n")
    f.write("ml system git \n")
    f.write("ml system git-lfs \n")
    f.write("ml curl/7.81.0 \n")
    f.write("ml system rclone \n")
    f.write("ml load matlab/R2022b \n")
    f.write("source /home/users/%s/miniconda3/bin/activate blp-instruments \n" %user)
    f.write("cd code/Sim \n")

    f.write('echo "Array Task ID: $SLURM_ARRAY_TASK_ID"\n')
    f.write(f"START=$((1 + {arrayjob} * ($SLURM_ARRAY_TASK_ID - 1)))\n")
    f.write(f"STOP=$((0 + {arrayjob} * $SLURM_ARRAY_TASK_ID))\n")
    f.write("matlab -nodisplay -nosplash -nodesktop \\\n")
    f.write("""    -r "for s=$START:$STOP; cal_distance(s, '%s', '%s', %s, ''); end; for s=$START:$STOP; cal_distance(s, '%s', '%s', %s, '_sh'); end; exit;"\n""" %
            (data_path, out_path, market, data_path, out_path, market))
    f.write("cd -")
    f.close()

    if dependency:
        cmd = "sbatch --depend=afterok:%s --parsable %s" % (dependency, job_file)
    else:
        cmd = "sbatch --parsable %s" % job_file

    print("Submitting job for %s" % jobname)
    status, jid = subprocess.getstatusoutput(cmd)

    if (status == 0):
        print("Job id for %s is %s" % (jobname, jid))
    else:
        print("Error submitting job for %s %s." % (jobname, status))
        print(subprocess.getstatusoutput(cmd))   

    return(jid)

def task_allocation_light(model, market):
    nodes    = "1"
    time     = "5:00:00"
    ntasks   = "1"
    nmem     = "20"
    array = ",".join([f"{i}-{i+4}" for i in range(1, 1002, 100)])
    arrayjob = 1
    return(nodes, ntasks, nmem, time, array, arrayjob)

def task_allocation_distance(market):
    nodes    = "1"
    time     = "30:00:00"
    ntasks   = "1"
    nmem     = "20"
    array = ",".join([f"{i}" for i in range(1, 1002, 100)])
    arrayjob = 1
    return(nodes, ntasks, nmem, time, array, arrayjob)

def task_allocation_full(model, par_rest, market):
    nodes = "1"
    ntasks = "1"
    time = "20:00:00"
    
    if market == 10000:
       nmem = "20"

    if par_rest == []:
       type = 0
    elif par_rest == [2]:
       type = 2
    elif par_rest == [3]:
       type = 3
    elif par_rest == [2,3]:
       type = 23

    array_map = {
        ('L', 0, 10000): 22,
        ('L_nm', 0, 10000): 22,
        ('RCL', 23, 10000): 110,
    }
    array_length = array_map.get((model, type, market), 1)
    array = f"1-{array_length}"
    
    if model == 'L':
       arrayjob = 50
    elif model == 'RCL':
       arrayjob = 10
    elif model == 'L_nm':
       arrayjob = 50
    return nodes, ntasks, nmem, time, array, arrayjob

def get_lists(model, price_iv, market):

    if market == 10000:
    
        if model == 'RCNL':
            par_rest_list = [[2,3]]
        
        elif model == 'RCL':
            par_rest_list = [[2,3]]
        
        else:
            par_rest_list = [[]]
    
    return par_rest_list


if __name__ == "__main__":
    Main()
